{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from CNNModel import CNN\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "from sklearn.datasets import fetch_openml\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "\n",
    "# 设置中文显示\n",
    "import matplotlib as mpl\n",
    "\n",
    "plt.rcParams[\"font.sans-serif\"] = [\"SimHei\"]  # 设置中文字体为黑体\n",
    "plt.rcParams[\"axes.unicode_minus\"] = False  # 解决负号显示问题\n",
    "\n",
    "\n",
    "def load_mnist_dataset(sample_size=1000):\n",
    "    \"\"\"\n",
    "    加载MNIST数据集\n",
    "\n",
    "    参数:\n",
    "    sample_size -- 如果提供，限制数据集大小为指定值\n",
    "\n",
    "    返回:\n",
    "    X_train, X_test, y_train, y_test -- 训练集和测试集\n",
    "    \"\"\"\n",
    "    print(\"正在加载MNIST数据集...\")\n",
    "    # 从OpenML加载MNIST数据集\n",
    "    X, y = fetch_openml(\"mnist_784\", version=1, return_X_y=True, as_frame=False)\n",
    "\n",
    "    # 将数据限制到指定大小\n",
    "    if sample_size and sample_size < len(X):\n",
    "        X = X[:sample_size]\n",
    "        y = y[:sample_size]\n",
    "\n",
    "    # 标准化数据\n",
    "    X = X / 255.0\n",
    "\n",
    "    # 将标签从字符转为整数\n",
    "    label_encoder = LabelEncoder()\n",
    "    y = label_encoder.fit_transform(y)\n",
    "\n",
    "    # 重塑数据为 (样本数, 通道数, 高度, 宽度)\n",
    "    X = X.reshape(-1, 1, 28, 28)\n",
    "\n",
    "    # 分割为训练集和测试集\n",
    "    X_train, X_test, y_train, y_test = train_test_split(\n",
    "        X, y, test_size=0.2, random_state=42\n",
    "    )\n",
    "\n",
    "    print(f\"数据加载完成：{X_train.shape[0]} 个训练样本，{X_test.shape[0]} 个测试样本\")\n",
    "    return X_train, X_test, y_train, y_test\n",
    "\n",
    "\n",
    "def evaluate_model(model, X_test, y_test, batch_size=32):\n",
    "    \"\"\"\n",
    "    评估模型在测试集上的表现\n",
    "\n",
    "    参数:\n",
    "    model -- CNN模型实例\n",
    "    X_test -- 测试数据\n",
    "    y_test -- 测试标签\n",
    "    batch_size -- 批量大小\n",
    "\n",
    "    返回:\n",
    "    accuracy -- 准确率\n",
    "    \"\"\"\n",
    "    num_samples = X_test.shape[0]\n",
    "    correct = 0\n",
    "\n",
    "    # 分批进行预测以处理大型数据集\n",
    "    for i in range(0, num_samples, batch_size):\n",
    "        end = min(i + batch_size, num_samples)\n",
    "        X_batch = X_test[i:end]\n",
    "        y_batch = y_test[i:end]\n",
    "\n",
    "        # 获取预测结果\n",
    "        predictions = model.predict(X_batch)\n",
    "\n",
    "        # 计算准确的预测数量\n",
    "        correct += np.sum(predictions == y_batch)\n",
    "\n",
    "    # 计算准确率\n",
    "    accuracy = correct / num_samples\n",
    "    return accuracy\n",
    "\n",
    "\n",
    "def visualize_results(losses, accuracies, save_path=None):\n",
    "    \"\"\"\n",
    "    可视化训练结果\n",
    "\n",
    "    参数:\n",
    "    losses -- 每个epoch的损失列表\n",
    "    accuracies -- 每个epoch的准确率列表\n",
    "    save_path -- 保存图像的路径，如果是None则显示图像\n",
    "    \"\"\"\n",
    "    epochs = range(1, len(losses) + 1)\n",
    "\n",
    "    plt.figure(figsize=(12, 5))\n",
    "\n",
    "    # 绘制损失曲线\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.plot(epochs, losses, \"b-\", label=\"训练损失\")\n",
    "    plt.title(\"训练损失\")\n",
    "    plt.xlabel(\"Epochs\")\n",
    "    plt.ylabel(\"损失\")\n",
    "    plt.legend()\n",
    "\n",
    "    # 绘制准确率曲线\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.plot(epochs, accuracies, \"r-\", label=\"测试准确率\")\n",
    "    plt.title(\"测试准确率\")\n",
    "    plt.xlabel(\"Epochs\")\n",
    "    plt.ylabel(\"准确率\")\n",
    "    plt.legend()\n",
    "\n",
    "    plt.tight_layout()\n",
    "\n",
    "    if save_path:\n",
    "        plt.savefig(save_path)\n",
    "    else:\n",
    "        plt.show()\n",
    "\n",
    "\n",
    "def visualize_predictions(model, X_test, y_test, num_samples=5):\n",
    "    \"\"\"\n",
    "    可视化预测结果\n",
    "\n",
    "    参数:\n",
    "    model -- CNN模型实例\n",
    "    X_test -- 测试数据\n",
    "    y_test -- 测试标签\n",
    "    num_samples -- 显示样本数量\n",
    "    \"\"\"\n",
    "    # 随机选择样本\n",
    "    indices = np.random.choice(len(X_test), num_samples, replace=False)\n",
    "\n",
    "    plt.figure(figsize=(15, 3))\n",
    "\n",
    "    for i, idx in enumerate(indices):\n",
    "        # 获取样本\n",
    "        image = X_test[idx, 0]  # 取出第一个通道\n",
    "        true_label = y_test[idx]\n",
    "\n",
    "        # 获取预测\n",
    "        prediction = model.predict(X_test[idx : idx + 1])[0]\n",
    "\n",
    "        # 显示图像和标签\n",
    "        plt.subplot(1, num_samples, i + 1)\n",
    "        plt.imshow(image, cmap=\"gray\")\n",
    "        plt.title(f\"预测: {prediction}\\n真实: {true_label}\")\n",
    "        plt.axis(\"off\")\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def main():\n",
    "    # 加载数据集（使用合适大小的样本以平衡速度和效果）\n",
    "    X_train, X_test, y_train, y_test = load_mnist_dataset(sample_size=1000)\n",
    "\n",
    "    # 设置训练参数\n",
    "    batch_size = 32  # 更小的批量大小\n",
    "    epochs = 3  # 减少轮次\n",
    "    learning_rate = 0.01\n",
    "\n",
    "    # 初始化CNN模型\n",
    "    input_shape = (1, 28, 28)  # 单通道，28x28像素\n",
    "    num_classes = 10  # MNIST有10个类别（数字0-9）\n",
    "    model = CNN(input_shape=input_shape, num_classes=num_classes)\n",
    "\n",
    "    print(\"开始训练模型...\")\n",
    "    start_time = time.time()\n",
    "\n",
    "    # 训练模型并记录损失和准确率\n",
    "    losses = []\n",
    "    accuracies = []\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        # 训练一个epoch\n",
    "        print(f\"Epoch {epoch+1}/{epochs} 训练中...\")\n",
    "        num_samples = X_train.shape[0]\n",
    "        total_loss = 0\n",
    "\n",
    "        # 打乱数据\n",
    "        indices = np.random.permutation(num_samples)\n",
    "        X_shuffled = X_train[indices]\n",
    "        y_shuffled = y_train[indices]\n",
    "\n",
    "        # 批量训练\n",
    "        for i in range(0, num_samples, batch_size):\n",
    "            end = min(i + batch_size, num_samples)\n",
    "            X_batch = X_shuffled[i:end]\n",
    "            y_batch = y_shuffled[i:end]\n",
    "\n",
    "            # 前向传播计算损失\n",
    "            loss = model.forward(X_batch, y_batch)\n",
    "            total_loss += loss\n",
    "\n",
    "            # 反向传播更新参数\n",
    "            model.backward(learning_rate)\n",
    "\n",
    "            # 只打印部分进度以减少输出\n",
    "            if (i // batch_size) % 5 == 0:\n",
    "                print(\n",
    "                    f\"Batch {i//batch_size+1}/{num_samples//batch_size+1}, Loss: {loss:.4f}\"\n",
    "                )\n",
    "\n",
    "        # 计算平均损失\n",
    "        avg_loss = total_loss / (num_samples // batch_size)\n",
    "        losses.append(avg_loss)\n",
    "\n",
    "        # 评估模型\n",
    "        accuracy = evaluate_model(model, X_test, y_test, batch_size)\n",
    "        accuracies.append(accuracy)\n",
    "\n",
    "        print(\n",
    "            f\"Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}, Accuracy: {accuracy:.4f}\"\n",
    "        )\n",
    "\n",
    "    # 计算总训练时间\n",
    "    training_time = time.time() - start_time\n",
    "    print(f\"训练完成！总时间: {training_time:.2f} 秒\")\n",
    "\n",
    "    # 最终评估\n",
    "    final_accuracy = evaluate_model(model, X_test, y_test, batch_size)\n",
    "    print(f\"最终测试准确率: {final_accuracy:.4f}\")\n",
    "\n",
    "    # 可视化训练结果\n",
    "    visualize_results(losses, accuracies)\n",
    "\n",
    "    # 可视化预测结果\n",
    "    visualize_predictions(model, X_test, y_test)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
